#include <glog/logging.h>
#include "eskf_filter.h"

namespace sad {

EskFilter::EskFilter() {
    Rwb = Eigen::Matrix3d::Identity();
    twb = Eigen::Vector3d::Zero();
}

bool EskFilter::Predict(const IMU& imu) {
    assert(imu.timestamp_ >= current_time_);

    double dt = imu.timestamp_ - current_time_;
    if (dt > (5 * options_.imu_dt_) || dt < 0) {
        // 时间间隔不对，可能是第一个IMU数据，没有历史信息
        LOG(INFO) << "skip this imu because dt_ = " << dt;
        current_time_ = imu.timestamp_;
        last_imu_ = imu;
        return false;
    }

    double dt2 = dt*dt;
    // Acc and gyro.
    const Eigen::Vector3d acc_unbias = 0.5 * (last_imu_.acce_ + imu.acce_) - state_.ba;
    const Eigen::Vector3d gyro_unbias = 0.5 * (last_imu_.gyro_ + imu.gyro_) - state_.bg;
    state_.pose += state_.velocity * dt + 0.5 * (state_.rot * acc_unbias + gravity_) * dt2;
    state_.velocity += (state_.rot * acc_unbias + gravity_) * dt;

    const Eigen::Vector3d delta_angle_axis = gyro_unbias * dt;
    
    /*LOG(INFO) << "last_imu_.gyro_ = " << last_imu_.gyro_.transpose();
    LOG(INFO) << "imu.gyro_ ====== " << imu.gyro_.transpose();
    LOG(INFO) << "gyro_unbias ==== " << gyro_unbias.transpose();
    LOG(INFO) << "dt ============= " << dt;
    //LOG(INFO) << "delta_angle_axis = " << delta_angle_axis.transpose();
    LOG(INFO) << "delta_angle_axis = " << delta_angle_axis.norm();
    //LOG(INFO) << "state_.pose = " << state_.pose.transpose();
    LOG(INFO) << "state_.pose = " << state_.pose.transpose();*/

    if (delta_angle_axis.norm() > 1e-12) {
        Eigen::Matrix3d delta_rot = Eigen::AngleAxisd(delta_angle_axis.norm(), delta_angle_axis.normalized()).toRotationMatrix();
        state_.rot *= delta_rot;
        //Eigen::Quaterniond delta_quat = Eigen::Quaterniond(1.0, delta_angle_axis[0]/2.0, delta_angle_axis[1]/2.0, delta_angle_axis[2]/2.0); 
        //state_.quat *= delta_quat;
        //state_.quat.normalized();
        //state_.rot = state_.quat.toRotationMatrix();

        //LOG(INFO) << "delta_angle_axis = " << delta_angle_axis.transpose();
        //LOG(INFO) << "state_.rot = " << state_.rot;
        //LOG(INFO) << "state_.quat = " << state_.quat.toRotationMatrix();
    }
   
    //state_.quat = state_.quat * Exp(delta_angle_axis);
    //state_.quat.normalized();
    //state_.rot = state_.quat.toRotationMatrix();
    //state_.quat = Eigen::Quaterniond(state_.rot);
    //state_.quat.normalized();
    //state_.rot = state_.quat.toRotationMatrix();
    //LOG(INFO) << "state_.pose 0= " << state_.pose.transpose();
    //LOG(INFO) << "state_.velocity 0= " << state_.velocity.transpose();
    //LOG(INFO) << "state_.rot 0= " << state_.rot;
    //LOG(INFO) << "state_.ba 0= " << state_.ba.transpose();
    //LOG(INFO) << "state_.bg 0= " << state_.bg.transpose();

    // Covariance of the error-state.   
    Eigen::Matrix<double, 18, 18> Fx = Eigen::Matrix<double, 18, 18>::Identity();
    Fx.block<3, 3>(0, 3)   = Eigen::Matrix3d::Identity() * dt;
    Fx.block<3, 3>(3, 6)   = -state_.rot * SkewMatrix(acc_unbias) * dt;
    Fx.block<3, 3>(3, 9)   = -state_.rot * dt;

    Fx.block<3, 3>(3, 15) = Eigen::Matrix3d::Identity() * dt; 
    if (/*delta_angle_axis.norm() > 1e-12*/true) {
        Fx.block<3, 3>(6, 6) = Eigen::AngleAxisd(delta_angle_axis.norm(), delta_angle_axis.normalized()).toRotationMatrix().transpose();
    } else {
        Fx.block<3, 3>(6, 6).setIdentity();
    }

    //Fx.block<3, 3>(6, 6) = Exp(-delta_angle_axis);

    Fx.block<3, 3>(6, 12)  = -Eigen::Matrix3d::Identity() * dt;

    Eigen::Matrix<double, 18, 12> Fi = Eigen::Matrix<double, 18, 12>::Zero();
    Fi.block<12, 12>(3, 0) = Eigen::Matrix<double, 12, 12>::Identity();

    Eigen::Matrix<double, 12, 12> Qi = Eigen::Matrix<double, 12, 12>::Zero();
    Qi.block<3, 3>(0, 0) = options_.acce_var_ * Eigen::Matrix3d::Identity();
    Qi.block<3, 3>(3, 3) = options_.gyro_var_ * Eigen::Matrix3d::Identity();
    Qi.block<3, 3>(6, 6) = options_.bias_acce_var_ * Eigen::Matrix3d::Identity();
    Qi.block<3, 3>(9, 9) = options_.bias_gyro_var_ * Eigen::Matrix3d::Identity();

    state_.cov = Fx * state_.cov * Fx.transpose() + Fi * Qi * Fi.transpose();

    state_.timestamp = imu.timestamp_;
    current_time_ = imu.timestamp_;
    last_imu_ = imu;
  
    return false;
}

bool EskFilter::ObserveGps(const GNSS& gnss) {
    //LOG(INFO) << "gnss.unix_time_ = " << gnss.unix_time_;

    /// GNSS 观测的修正
    assert(gnss.unix_time_ >= current_time_);
    gnss_ = gnss;
    //LOG(INFO) << "first_gnss_ = " << first_gnss_;

    if (first_gnss_) {
        state_.rot = gnss.utm_pose_.so3().matrix();
        state_.pose = gnss.utm_pose_.translation();
        first_gnss_ = false;
        current_time_ = gnss.unix_time_;
        return true;
    }

    assert(gnss.heading_valid_);
    //LOG(INFO) << "gnss.heading_valid_ = " << gnss.heading_valid_;
    ObserveSE3(gnss.utm_pose_, options_.gnss_pos_noise_, options_.gnss_ang_noise_);
    //ObservePose(gnss.utm_pose_, options_.gnss_pos_noise_, options_.gnss_ang_noise_);
    current_time_ = gnss.unix_time_;
}

bool EskFilter::ObserveSE3(const SE3& pose, double trans_noise, double ang_noise) {
    Eigen::Matrix3d R = pose.so3().matrix();
    Eigen::Vector3d axis_angle = Log(state_.rot);

    Eigen::Matrix<double, 3, 1> t = pose.translation();

    Eigen::Matrix<double, 6, 18> H;
    H.setZero();
    H.block<3,3>(0,0) = Eigen::Matrix3d::Identity();
    H.block<3,3>(0,6) = Eigen::Matrix3d::Identity();

    Eigen::Matrix<double, 3, 1> lat_lon_alt_std(gnss_.lat_lon_alt_std_[0], gnss_.lat_lon_alt_std_[1], gnss_.lat_lon_alt_std_[2]);

    Eigen::Matrix<double, 6, 1> noise_vec;
    //noise_vec << trans_noise, trans_noise, trans_noise;
    noise_vec << lat_lon_alt_std[0], lat_lon_alt_std[1], lat_lon_alt_std[2], ang_noise, ang_noise, ang_noise;

    Eigen::Matrix<double, 6, 6> V = noise_vec.asDiagonal();
    Eigen::Matrix<double, 18, 6> K = state_.cov * H.transpose() * (H * state_.cov * H.transpose() + V).inverse();

    // 更新x和cov
    Eigen::Matrix<double, 6, 1> innov = Eigen::Matrix<double, 6, 1>::Zero();
    innov.template head<3>() = (pose.translation() - state_.pose);          // 平移部分
    innov.template tail<3>() = Log(state_.rot.inverse() * R);  // 旋转部分(3.67)

    Eigen::Matrix<double, 18, 1> delta_x = K * innov;
    // state update
    state_.pose += delta_x.block<3, 1>(0, 0);
    state_.velocity += delta_x.block<3, 1>(3, 0);
    state_.ba += delta_x.block<3, 1>(9, 0);
    state_.bg += delta_x.block<3, 1>(12, 0);

    gravity_ +=  delta_x.block<3,1>(15,0);

    if (/*delta_x.block<3, 1>(6, 0).norm() > 1e-12*/true) {
        state_.rot *= Eigen::AngleAxisd(delta_x.block<3, 1>(6, 0).norm(), delta_x.block<3, 1>(6, 0).normalized()).toRotationMatrix();
    }

    //state_.quat = Eigen::Quaterniond(state_.rot);
    //state_.quat.normalized();
    //state_.rot = state_.quat.toRotationMatrix();

    LOG(INFO) << "state_.pose = " << state_.pose.transpose();

    // Covarance.
    state_.cov = (Eigen::Matrix<double, 18, 18>::Identity() - K * H) * state_.cov;

    return true;
}

bool EskFilter::ObservePose(const SE3& pose, double trans_noise, double ang_noise) {
    Eigen::Matrix<double, 3, 1> t = pose.translation();
  
    Eigen::Matrix<double, 3, 18> H;
    H.setZero();
    H.block<3,3>(0,0) = Eigen::Matrix3d::Identity();
   
    Eigen::Matrix<double, 3, 1> residual;
    residual.block<3, 1>(0, 0) = t - state_.pose; 

    //LOG(INFO) << "state_.pose = " << state_.pose.transpose();
    //LOG(INFO) << "t =========== " << t.transpose();
    //LOG(INFO) << "residual = " << residual.transpose();
    double gnss_info_x_scale = 1.0;
    double gnss_info_y_scale = 1.0;
    double gnss_info_z_scale = 1.0;
    Vec3d lat_lon_alt_std(gnss_info_x_scale * gnss_.lat_lon_alt_std_[0], gnss_info_y_scale * gnss_.lat_lon_alt_std_[1], gnss_info_z_scale * gnss_.lat_lon_alt_std_[2]);

    Eigen::Matrix<double, 3, 1> noise_vec;
    //noise_vec << trans_noise, trans_noise, trans_noise;
    noise_vec << lat_lon_alt_std[0], lat_lon_alt_std[1], lat_lon_alt_std[2];

    //noise_vec << lat_lon_alt_std[0], lat_lon_alt_std[1], lat_lon_alt_std[2], ang_noise, ang_noise, ang_noise, ang_noise;

    Eigen::Matrix<double, 3, 3> V = noise_vec.asDiagonal();

    // EKF.
    Eigen::Matrix<double, 18, 18> P = state_.cov;
    const Eigen::Matrix<double, 3, 3> S = H * P * H.transpose() + V;
    const Eigen::Matrix<double, 3, 3> S_inv = S.llt().solve(Eigen::MatrixXd::Identity(S.rows(), S.cols()));
    Eigen::Matrix<double, 18, 3> K = P * H.transpose() * S_inv;
    //Eigen::Matrix<double, 15, 3> K = P * H.transpose() * (H * P * H.transpose() + V).inverse();
    Eigen::Matrix<double, 18, 1> delta_x = K * residual;
    
    // state update
    state_.pose += delta_x.block<3, 1>(0, 0);
    state_.velocity += delta_x.block<3, 1>(3, 0);
    state_.ba += delta_x.block<3, 1>(9, 0);
    state_.bg += delta_x.block<3, 1>(12, 0);

    gravity_ +=  delta_x.block<3,1>(15,0);
#if 1
    if (delta_x.block<3, 1>(6, 0).norm() > 1e-12) {
        state_.rot *= Eigen::AngleAxisd(delta_x.block<3, 1>(6, 0).norm(), delta_x.block<3, 1>(6, 0).normalized()).toRotationMatrix();
        //Eigen::Vector3d delta_theta = delta_x.block<3, 1>(6, 0);
        //state_.quat *= Eigen::Quaterniond(1.0, delta_theta[0]/2.0, delta_theta[1]/2.0, delta_theta[2]/2.0);
 
        //LOG(INFO) << "state_.rot = " << state_.rot;
    }
    //state_.quat = Eigen::Quaterniond(state_.rot);

    //state_.quat.normalized();
    //state_.rot = state_.quat.toRotationMatrix();
#else

    Eigen::Vector3d error_state_rot = delta_x.block<3,1>(6,0);
    double error_state_rot_norm = error_state_rot.norm();
    if(error_state_rot_norm > 1e-12){
        error_state_rot /= error_state_rot_norm;
        error_state_rot *= std::sin(error_state_rot_norm/2.0);
        Eigen::Quaterniond error_state_q_(std::cos(error_state_rot_norm/2),error_state_rot[0],error_state_rot[1],error_state_rot[2]);
        state_.quat = state_.quat * error_state_q_;
    }
    state_.quat.normalize();

    LOG(INFO) << "error_state_rot = " << error_state_rot.transpose();

#endif
    LOG(INFO) << "state_.pose = " << state_.pose.transpose();
    LOG(INFO) << "state_.velocity = " << state_.velocity.transpose();
    LOG(INFO) << "state_.rot = " << std::endl << state_.quat.toRotationMatrix();
    LOG(INFO) << "state_.ba = " << state_.ba.transpose();
    LOG(INFO) << "state_.bg = " << state_.bg.transpose();
    // Covarance.
    //Eigen::Matrix<double, 15, 15> KH = Eigen::Matrix<double, 15, 15>::Identity() - K * H;
    //state_.cov = KH * P * KH.transpose() + K * V * K.transpose();
    
    state_.cov = (Eigen::Matrix<double, 18, 18>::Identity() - K * H) * P;

    return true;
}

bool EskFilter::ObserveQuaternion(const SE3& pose, double trans_noise, double ang_noise) {
    /*Eigen::Matrix3d R = pose.so3().matrix();
    Eigen::Vector3d axis_angle = Log(state_.rot);

    Eigen::Matrix<double, 3, 1> t = pose.translation();
    Eigen::Matrix<double, 4, 3> Qtheta;

    state_.quat = Eigen::Quaterniond(state_.rot);
    state_.quat.normalized();

    Qtheta(0,0) = -state_.quat.x();
    Qtheta(0,1) = -state_.quat.y();
    Qtheta(0,2) = -state_.quat.z();
    Qtheta(1,0) = state_.quat.w();
    Qtheta(1,1) = -state_.quat.z();
    Qtheta(1,2) = state_.quat.y();
    Qtheta(2,0) = state_.quat.z();
    Qtheta(2,1) = state_.quat.w();
    Qtheta(2,2) = -state_.quat.x();
    Qtheta(3,0) = -state_.quat.y();
    Qtheta(3,1) = state_.quat.x();
    Qtheta(3,2) = state_.quat.w();

    Qtheta *= 0.5;
   
    Eigen::Matrix<double, 7, 15> H;
    H.setZero();
    H.block<3,3>(0,0) = Eigen::Matrix3d::Identity();
 
    H.block<4,3>(3,6) = Qtheta;

    Eigen::Quaterniond temp_quat = state_.quat * Eigen::Quaterniond(R).conjugate();

    temp_quat.normalized();

    Eigen::Matrix<double, 4, 1> quat;
    quat(0, 0) = temp_quat.w();
    quat(1, 0) = temp_quat.x();
    quat(2, 0) = temp_quat.y();
    quat(3, 0) = temp_quat.z();    
    Eigen::Matrix<double, 7, 1> residual;
    residual.block<3, 1>(0, 0) = state_.pose - t; 
    residual.block<4, 1>(3, 0) = quat;

    LOG(INFO) << "state_.pose = " << state_.pose.transpose();
    LOG(INFO) << "t =========== " << t.transpose();
    LOG(INFO) << "residual = " << residual.transpose();

    Vec3d lat_lon_alt_std(gnss_.lat_lon_alt_std_[0], gnss_.lat_lon_alt_std_[1], gnss_.lat_lon_alt_std_[2]);

    Eigen::Matrix<double, 7, 1> noise_vec;
    //noise_vec << trans_noise, trans_noise, trans_noise;
    noise_vec << lat_lon_alt_std[0], lat_lon_alt_std[1], lat_lon_alt_std[2], ang_noise, ang_noise, ang_noise, ang_noise;

    Eigen::Matrix<double, 7, 7> V = noise_vec.asDiagonal();

    // EKF.
    Eigen::Matrix<double, 15, 15> P = state_.cov;
    Eigen::Matrix<double, 15, 7> K = P * H.transpose() * (H * P * H.transpose() + V).inverse();
    Eigen::Matrix<double, 15, 1> delta_x = K * residual;
    
    // state update
    state_.pose += delta_x.block<3, 1>(0, 0);
    state_.velocity += delta_x.block<3, 1>(3, 0);
    state_.ba += delta_x.block<3, 1>(9, 0);
    state_.bg += delta_x.block<3, 1>(12, 0);

    if (delta_x.block<3, 1>(6, 0).norm() > 1e-12) {
        //state_.rot *= Eigen::AngleAxisd(delta_x.block<3, 1>(6, 0).norm(), delta_x.block<3, 1>(6, 0).normalized()).toRotationMatrix();
        //Eigen::Vector3d delta_theta = delta_x.block<3, 1>(6, 0);
        //state_.quat *= Eigen::Quaterniond(1.0, delta_theta[0]/2.0, delta_theta[1]/2.0, delta_theta[2]/2.0);
        Eigen::Vector4d delta_q = delta_x.block<4, 1>(6, 0);

        state_.quat *= Eigen::Quaterniond(delta_q[3], delta_q[0], delta_q[1], delta_q[2]);

        state_.quat.normalized();

        state_.rot = state_.quat.toRotationMatrix();
        //LOG(INFO) << "state_.rot = " << state_.rot;
    }

    LOG(INFO) << "state_.pose = " << state_.pose.transpose();

    // Covarance.
    //Eigen::Matrix<double, 15, 15> KH = Eigen::Matrix<double, 15, 15>::Identity() - K * H;
    //state_.cov = KH * P * KH.transpose() + K * V * K.transpose();
    
    state_.cov = (Eigen::Matrix<double, 15, 15>::Identity() - K * H) * P;*/

    return true;
}

bool EskFilter::ObserveWheelSpeed(const Odom& odom) {
    assert(odom.timestamp_ >= current_time_);

    options_.odom_span_ = odom.timestamp_ - last_odom_.timestamp_;
    if(options_.odom_span_ > 1.0 || options_.odom_span_ < 1.0e-5) {
        options_.odom_span_ = 0.01;
    }     
    double velo_l = options_.left_wheel_radius_ * (odom.left_pulse_) / options_.circle_pulse_ * 2 * M_PI / options_.odom_span_;
    double velo_r = options_.right_wheel_radius_ * (odom.right_pulse_) / options_.circle_pulse_ * 2 * M_PI / options_.odom_span_;   
    double average_vel = 0.5 * (velo_l + velo_r);
    
    last_odom_ = odom;
    return true;
}

Eigen::Matrix3d EskFilter::Exp(const Eigen::Vector3d &axis_angle) {
    double norm = axis_angle.norm();
    if (norm > 1e-12) {
        Eigen::Vector3d r = axis_angle / norm; 
        Eigen::Matrix3d skew_axis = SkewMatrix(r); 
        double theta = norm; 
        return Eigen::Matrix3d::Identity() + std::sin(theta) * skew_axis + (1.0 - std::cos(theta)) * skew_axis * skew_axis;
    }

    return Eigen::Matrix3d::Identity();
}

Eigen::Vector3d EskFilter::Log(const Eigen::Matrix3d& R) {
#if 0
    LOG(INFO) << "R= " << R;

    Eigen::Vector3d axis_angle;
    double trace = R(0,0)+R(1,1)+R(2,2);
    double x = (R.trace()-1.0)/2.0;
    if(x > 1.0) {
        x = 1.0;
    } 
    if(x < -1.0) {
        x = -1.0;
    }

    double angle = acos(x);

    LOG(INFO) << "trace= " << trace;
    LOG(INFO) << "R.trace()= " << R.trace();
    LOG(INFO) << "(R.trace()-1.0)/2.0= " << (R.trace()-1.0)/2.0;

    Eigen::Matrix3d rot_minus_rot_trans = R - R.transpose();
    LOG(INFO) << "angle 0= " << angle;
    //LOG(INFO) << "rot_minus_rot_trans = " << rot_minus_rot_trans;
    double wx = (rot_minus_rot_trans(2,1) - rot_minus_rot_trans(1,2))/2.0;
    double wy = (rot_minus_rot_trans(0,2) - rot_minus_rot_trans(2,0))/2.0;
    double wz = (rot_minus_rot_trans(1,0) - rot_minus_rot_trans(0,1))/2.0;

    Eigen::Vector3d vee_axis(wx,wy,wz);
    axis_angle = vee_axis / (2*sin(angle));
    LOG(INFO) << "axis_angle 0= " << axis_angle.transpose();
#endif
    Eigen::AngleAxisd rotation_vector(R);

    double angle = rotation_vector.angle();
    Eigen::Vector3d axis_angle = rotation_vector.axis();

    //LOG(INFO) << "angle 1= " << angle;
    //LOG(INFO) << "axis_angle 1= " << axis_angle.transpose();

    return angle*axis_angle;
}

}    